from scipy.interpolate import CubicSpline
import numpy as np


class Path:
    wayPoints: np.ndarray  # Shape[n,2]
    interplotPoints: np.ndarray  # Shape[n,4] theta,vec,x,y

    def __init__(self, path: np.ndarray):  # Shape[n,2]
        self.wayPoints = path

    def interplot(self, steps: int) -> None:
        index = np.arange(0, self.wayPoints.shape[0])
        cs = CubicSpline(index, self.wayPoints)
        positions = cs(np.linspace(0, len(index) - 1, steps))
        xd = np.diff(positions[:, 0])
        yd = np.diff(positions[:, 1])
        angle = np.arctan(xd, yd).reshape(-1, 1)
        vec = np.sqrt(xd**2 + yd**2).reshape(-1, 1)
        self.interplotPoints = np.hstack([angle, vec, positions[0 : steps - 1, :]])

    def getWayPoints(self) -> np.ndarray:
        return self.wayPoints

    def getInterplotPoints(self) -> np.ndarray:
        return self.interplotPoints

    def getNearestPointIndex(self, x: float, y: float) -> int:
        distance = (self.interplotPoints[:, 2] - x) ** 2 + (
            self.interplotPoints[:, 3] - y
        ) ** 2
        index = distance.argmin()
        return int(index)
